import numpy as np
import matplotlib.pyplot as plt
import matplotlib as mpl
from sklearn.linear_model import LinearRegression
from sklearn.metrics import r2_score

# 设置全局字体配置
plt.rcParams['font.size'] = 10
plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei', 'Arial Unicode MS']  # 多字体备选
plt.rcParams['axes.unicode_minus'] = False  # 正确显示负号
plt.rcParams['text.usetex'] = False  # 禁用LaTeX渲染

# 创建支持中文和数学符号的字体属性
chinese_font = mpl.font_manager.FontProperties(fname=mpl.font_manager.findfont('SimHei'))
math_font = mpl.font_manager.FontProperties(fname=mpl.font_manager.findfont('DejaVu Sans'))

# 设置随机种子保证可复现性
np.random.seed(42)

# 实际梯度范数数据 (范围0.1-1.0)
actual_norms = np.linspace(0.1, 1.0, 100)

# 预测梯度范数 (加入预测误差)
predicted_norms = 0.85 * actual_norms + 0.15 * np.random.normal(0, 0.08, 100)
predicted_norms = np.clip(predicted_norms, 0.05, 1.05)  # 限制在合理范围

# 计算R²拟合优度
r2 = r2_score(actual_norms, predicted_norms)

# 创建回归线模型
model = LinearRegression()
model.fit(actual_norms.reshape(-1, 1), predicted_norms)
regression_line = model.predict(actual_norms.reshape(-1, 1))

# 创建图形
fig, ax = plt.subplots(figsize=(8, 6), dpi=100)

# 绘制散点图
scatter = ax.scatter(actual_norms, predicted_norms, 
                     c='#1f77b4', alpha=0.7, 
                     edgecolors='w', linewidth=0.5,
                     label=f'客户端数据点 (R$^2$={r2:.3f})')  # 使用LaTeX语法显示上标

# 绘制回归线
ax.plot(actual_norms, regression_line, 
        'r--', linewidth=1.5, 
        label='线性回归拟合')

# 绘制理想预测线 (y=x)
ax.plot([0, 1], [0, 1], 
        'g-', linewidth=1, 
        label='理想预测线')

# 设置坐标轴范围
ax.set_xlim(0.05, 1.05)
ax.set_ylim(0.05, 1.05)

# 添加标题和标签 - 使用中文字体
ax.set_title('图4：梯度范数预测准确性验证', fontsize=14, fontweight='bold', fontproperties=chinese_font)
ax.set_xlabel('实际梯度范数 ||∇W||', fontsize=12, fontproperties=chinese_font)
ax.set_ylabel('预测梯度范数 ||∇Ŵ||', fontsize=12, fontproperties=chinese_font)

# 添加网格
ax.grid(True, linestyle='--', alpha=0.7)

# 添加图例 - 使用混合字体
legend = ax.legend(loc='upper left', frameon=True, prop=math_font)
for text in legend.get_texts():
    # 检测文本中是否包含中文
    if any('\u4e00' <= char <= '\u9fff' for char in text.get_text()):
        text.set_fontproperties(chinese_font)
    else:
        text.set_fontproperties(math_font)
        
# 设置图例背景透明度
legend.get_frame().set_alpha(0.95)

# 添加技术标注 - 使用混合字体
fig.text(0.5, 0.01, 
         f"GP-AdaFL框架 | 数据集: MNIST+CIFAR-10 | 预测窗口大小: 5轮 | R$^2$={r2:.3f}", 
         ha="center", fontsize=9, style='italic', fontproperties=chinese_font)

# 显示并保存
plt.tight_layout()
plt.savefig('gradient_prediction_accuracy.png', bbox_inches='tight', dpi=300)
plt.show()
